import numpy as np
import matplotlib.pyplot as plt
from scipy.special import comb


class BezierCurve:
    def __init__(self, control_points):
        self.control_points = np.array(control_points)
        self.n = len(control_points) - 1  # 曲线阶数

    def evaluate(self, t):
        """计算贝塞尔曲线上的点"""
        n = self.n
        point = np.zeros(2)

        for i in range(n + 1):
            # 伯恩斯坦多项式
            bernstein = comb(n, i) * (1 - t) ** (n - i) * t ** i
            point += bernstein * self.control_points[i]

        return point

    def get_curve_points(self, num_points=100):
        """获取曲线上的多个点"""
        t_values = np.linspace(0, 1, num_points)
        curve_points = np.array([self.evaluate(t) for t in t_values])
        return curve_points

    def plot(self, num_points=100, show_control_points=True):
        """绘制贝塞尔曲线"""
        curve_points = self.get_curve_points(num_points)

        plt.figure(figsize=(10, 8))
        plt.plot(curve_points[:, 0], curve_points[:, 1], 'b-', linewidth=2, label='贝塞尔曲线')

        if show_control_points:
            plt.plot(self.control_points[:, 0], self.control_points[:, 1],
                     'ro--', linewidth=1, markersize=6, label='控制点')

        plt.grid(True, alpha=0.3)
        plt.axis('equal')
        plt.legend()
        plt.xlabel('X')
        plt.ylabel('Y')
        plt.title(f'{self.n}阶贝塞尔曲线 - Omega形转弯')
        plt.show()


def create_omega_turn_4th_order(width=10, height=8):
    """创建四阶贝塞尔曲线的Omega形转弯"""
    # 控制点定义 - 形成Omega形状
    control_points = np.array([
        [0, 0],  # P0: 起点
        [width * 0.3, height * 0.8],  # P1: 控制向上弯曲
        [width * 0.5, height * 1.2],  # P2: 顶部控制点
        [width * 0.7, height * 0.8],  # P3: 控制向下弯曲
        [width, 0]  # P4: 终点
    ])
    return BezierCurve(control_points)


def create_omega_turn_5th_order(width=10, height=8):
    """创建五阶贝塞尔曲线的Omega形转弯"""
    # 控制点定义 - 形成更平滑的Omega形状
    control_points = np.array([
        [0, 0],  # P0: 起点
        [width * 0.2, height * 0.6],  # P1: 第一个控制点
        [width * 0.4, height * 1.0],  # P2: 第二个控制点
        [width * 0.6, height * 1.0],  # P3: 第三个控制点
        [width * 0.8, height * 0.6],  # P4: 第四个控制点
        [width, 0]  # P5: 终点
    ])
    return BezierCurve(control_points)


def create_pear_shape_turn(width=10, height=8):
    """创建梨形（灯泡形）转弯曲线"""
    # 控制点定义 - 形成梨形/灯泡形
    control_points = np.array([
        [0, 0],  # P0: 起点
        [width * 0.1, height * 0.7],  # P1: 控制初始弯曲
        [width * 0.3, height * 1.5],  # P2: 形成灯泡的宽大部分
        [width * 0.7, height * 1.2],  # P3: 控制顶部形状
        [width * 0.9, height * 0.4],  # P4: 控制收尾弯曲
        [width, 0]  # P5: 终点
    ])
    return BezierCurve(control_points)


def calculate_curvature(bezier_curve, t, delta=1e-5):
    """计算曲线上某点的曲率"""
    # 计算一阶导数
    point1 = bezier_curve.evaluate(t - delta)
    point2 = bezier_curve.evaluate(t + delta)
    dx_dt = (point2[0] - point1[0]) / (2 * delta)
    dy_dt = (point2[1] - point1[1]) / (2 * delta)

    # 计算二阶导数
    point0 = bezier_curve.evaluate(t - 2 * delta)
    point3 = bezier_curve.evaluate(t + 2 * delta)
    d2x_dt2 = (point3[0] - 2 * point2[0] + point1[0]) / (4 * delta ** 2)
    d2y_dt2 = (point3[1] - 2 * point2[1] + point1[1]) / (4 * delta ** 2)

    # 计算曲率
    numerator = abs(dx_dt * d2y_dt2 - dy_dt * d2x_dt2)
    denominator = (dx_dt ** 2 + dy_dt ** 2) ** 1.5

    if denominator == 0:
        return 0
    return numerator / denominator


# 示例使用
if __name__ == "__main__":
    # 创建不同类型的转弯曲线
    omega_4th = create_omega_turn_4th_order()
    omega_5th = create_omega_turn_5th_order()
    pear_shape = create_pear_shape_turn()

    # 绘制曲线
    print("四阶Omega形转弯曲线:")
    omega_4th.plot()

    print("五阶Omega形转弯曲线:")
    omega_5th.plot()

    print("梨形/灯泡形转弯曲线:")
    pear_shape.plot()

    # 计算曲率示例
    t_values = np.linspace(0, 1, 20)
    curvatures = [calculate_curvature(omega_5th, t) for t in t_values]

    plt.figure(figsize=(10, 6))
    plt.plot(t_values, curvatures, 'b-', linewidth=2)
    plt.grid(True, alpha=0.3)
    plt.xlabel('参数 t')
    plt.ylabel('曲率')
    plt.title('五阶贝塞尔曲线曲率变化')
    plt.show()

    # 输出曲线信息
    curve_points = omega_5th.get_curve_points()
    print(f"曲线长度: {len(curve_points)} 个点")
    print(f"起点: {curve_points[0]}")
    print(f"终点: {curve_points[-1]}")